import asyncio
import websockets
import urllib.parse
import json
from sparkai.llm.llm import ChatSparkLLM, ChunkPrintHandler
from sparkai.core.messages import ChatMessage
import os

# 星火认知大模型调用秘钥信息
SPARKAI_URL = 'wss://spark-api.xf-yun.com/v3.1/chat'
SPARKAI_APP_ID = '1d38b17a'
SPARKAI_API_SECRET = 'MTJjOGQ5YzE1ZWZjYjJkOWUxZTM4Mzcx'
SPARKAI_API_KEY = 'b33d686049c006fd12a7e9b345ec0409'
# 星火认知大模型Spark Pro的domain值
SPARKAI_DOMAIN = 'generalv3'

os.environ["http_proxy"] = "http://127.0.0.1:7890"
os.environ["https_proxy"] = "http://127.0.0.1:7890"

spark = ChatSparkLLM(
    spark_api_url=SPARKAI_URL,
    spark_app_id=SPARKAI_APP_ID,
    spark_api_key=SPARKAI_API_KEY,
    spark_api_secret=SPARKAI_API_SECRET,
    spark_llm_domain=SPARKAI_DOMAIN,
    streaming=False,
)

identify_prompt = """
任务：现在有一个错题本的APP，错题本的页面如下， 现在需要你根据下面的用户输入，判断是否需要跳转某一页面，如果用户想跳转到某一页面，麻烦你对其正确输出,，当然你的输出需要满足一定JSON格式。而且，你不能添加其他的废话，只需按照格式输出即可
首先，你得需要判断用户输入的话，是否是想要进行页面的跳转。如果需要进行跳转，那么输出的json中的“is_skip”属性是yes，然后“url”属性为对应需要跳转页面的URL。如果不需要进行跳转，则"is_skip"属性为false。
所以，你最重要的一点是，判断，用户输入的话，是否是想要跳转到某一页面。
页面如下：
首页： "pages/HomePage/HomePage"
首页-学习目标： "pages/HomePage/TargetPage"
首页-智能巩固： "pages/HomePage/AIStudyPage"
首页-我的概览： "pages/HomePage/MyOverviewPage"
首页-智能练习： "pages/HomePage/GenExercisesPage"
错题本： "pages/QuesPage/QuesPage"
学情分析： "pages/GrowPage/GrowPage"
我的： "pages/MinePage/MinePage"
拍照上传错题： "pages/newques/CameraPage/CameraPage"
我的-使用说明： "pages/minedeep/MineInstructionPage/MineInstructionPage"
我的-问题反馈： "pages/minedeep/MineProblemPage/MineProblemPage"
我的-用户协议： "pages/minedeep/MineAgreementPage/MineAgreementPage"
输出格式如下：
{
 "is_skip" : "yes",
 "url":"pages/MinePage/MinePage"
}
其中，is_skip的取值只有yes和no两种取值，没有其他取值。“url”是在上述中需要跳转的页面的URL。
用户输入：{}
"""

async def echo(websocket, path):
    history_chat = []
    system_input = {
        "role": "system",
        "content": "You are a helpful assistant. Answer all questions to the best of your ability."
    }
    history_chat.append(system_input)
    async for message in websocket:
        i_prompt = identify_prompt.format(str(message))
        print(i_prompt)
        messages = [ChatMessage(role="user", content=i_prompt)]
        handler = ChunkPrintHandler()
        identify_res = spark.generate([messages], callbacks=[handler])
        identify_res_json = json.loads(identify_res)
        user_input = {
            "role": "user",
            "content": message
        }
        history_chat.append(user_input)
        print(history_chat)
        response = spark.generate([history_chat], callbacks=[handler])
        print("=======================回复结果如下================")
        print(identify_res)
        if identify_res_json['is_skip'] == "yes":
            print("需要进行页面跳转")
            print(f"url为{identify_res_json['url']}")
            await websocket.send("12138")
            await asyncio.sleep(0.1)  # 模拟每秒发送一个数据
            await websocket.send(identify_res_json['url'])
            ai_chat = {
                'role': 'assistant',
                'content': f"已帮助您跳转到{identify_res_json['url']}页面,您可以在这里完成您的操作"
            }
            history_chat.append(ai_chat)
        else:
            print("不需要进行页面跳转")
            await websocket.send("2580")
            reply = ""
            for chunk in response:
                content = str(chunk.choices[0].delta.content)
                if content != "None":
                    reply = reply + str(content)
                await websocket.send(content)
                await asyncio.sleep(0.1)  # 模拟每秒发送一个数据
            print(reply)
            ai_chat = {
                "role": "assistant",
                "content": reply
            }
            history_chat.append(ai_chat)

start_server = websockets.serve(echo, "localhost", 8765)
asyncio.get_event_loop().run_until_complete(start_server)
asyncio.get_event_loop().run_forever()
